import numpy as np
import geatpy as ea
import subprocess


# Utility functions


def find_extremes(ds):
    """Given a disp series, return the index list of extreme disps."""
    exts = []
    for i in range(1, len(ds)-1):
        if ds[i] > ds[i-1] and ds[i] > ds[i+1] or ds[i] < ds[i-1] and ds[i] < ds[i+1]:
            exts.append(i)
    exts.append(len(ds) - 1)
    return exts


def resample(ds, fs, dmax):
    """Resample the disp and force series.
    Returns:
       `ds`: resampled disps
       `fs`: resampled forces
    """
    rds, rfs = [], []
    i = 0
    dhist = 0.0
    exts = find_extremes(ds)
    while i < len(ds):
        if i == 0 or i in exts:
            rds.append(ds[i])
            rfs.append(fs[i])
            dhist = ds[i]
        else:
            if abs(ds[i] - dhist) > dmax:
                rds.append(ds[i-1])
                rfs.append(fs[i-1])
                dhist = ds[i-1]
        i += 1
    rds.append(ds[-1])
    rfs.append(fs[-1])
    return rds, rfs


def mesh_protocol(protocol, dt):
    """Given a protocol, mesh it into displacement points."""
    p = [0] + protocol
    res = []
    for i in range(1, len(p)):
        res.append(np.linspace(p[i-1], p[i], int(abs(p[i]-p[i-1]) / dt), endpoint=False))
    res.append(protocol[-1])
    return np.hstack(res).tolist()


def generate_script(template, paramlist, disps):
    """Given a list of material parameters and a list of disps, generate OpenSees script."""
    matstring = ""
    dofs = 1
    if template == "doweltype_exp":
        for i, params in enumerate(paramlist):
            matstring += "uniaxialMaterial DowelType {0} {1} -exponential {2}\n".format(
                i,
                " ".join([str(p) for p in list(params[:11])]),
                " ".join([str(p) for p in list(params[11:])])
            )
    elif template == "pinching4_eng":
        for i, params in enumerate(paramlist):
            matstring += "uniaxialMaterial Pinching4 {0} {1} energy\n".format(
                i,
                " ".join([str(p) for p in list(params)])
            )
    elif template == "doweltype_pw":
        for i, params in enumerate(paramlist):
            matstring += "uniaxialMaterial DowelType {0} {1} -piecewise {2}\n".format(
                i,
                " ".join([str(p) for p in list(params[:11])]),
                " ".join([str(p) for p in list(params[11:])])
            )
    elif template == "doweltype_bez":
        for i, params in enumerate(paramlist):
            matstring += "uniaxialMaterial DowelType {0} {1} -bezier {2}\n".format(
                i,
                " ".join([str(p) for p in list(params[:11])]),
                " ".join([str(p) for p in list(params[11:])])
            )
    elif template == "steel_doweltype_exp":
        for i, params in enumerate(paramlist):
            matstring += """uniaxialMaterial DowelType {0} {1} -exponential {2}
                uniaxialMaterial Steel01 {3} {4} {5} {6}
                uniaxialMaterial Series {7} {8} {9}\n""".format(
                100000 + i,
                " ".join([str(p) for p in list(params[3:14])]),
                " ".join([str(p) for p in list(params[14:])]),
                200000 + i,
                params[0],
                params[1],
                params[2],
                i,
                100000 + i,
                200000 + i
            )
    elif template == "steel_doweltype_bez":
        for i, params in enumerate(paramlist):
            matstring += """uniaxialMaterial DowelType {0} {1} -bezier {2}
                uniaxialMaterial Steel01 {3} {4} {5} {6}
                uniaxialMaterial Series {7} {8} {9}\n""".format(
                100000 + i,
                " ".join([str(p) for p in list(params[3:14])]),
                " ".join([str(p) for p in list(params[14:])]),
                200000 + i,
                params[0],
                params[1],
                params[2],
                i,
                100000 + i,
                200000 + i
            )
    elif template == "steel_hdn_doweltype_exp":
        for i, params in enumerate(paramlist):
            matstring += """uniaxialMaterial DowelType {0} {1} -exponential {2}
                uniaxialMaterial Steel01 {3} {4}
                uniaxialMaterial Series {5} {6} {7}\n""".format(
                100000 + i,
                " ".join([str(p) for p in list(params[7:18])]),
                " ".join([str(p) for p in list(params[18:])]),
                200000 + i,
                " ".join([str(p) for p in list(params[:7])]),
                i,
                100000 + i,
                200000 + i
            )
    elif template == "hook_pinching4_eng":
        # params: stiffness, activation, slot, pyratio, pinching4...
        for i, params in enumerate(paramlist):
            matstring += """uniaxialMaterial Pinching4 {0} {1} energy
                uniaxialMaterial Steel01 {2} {3} {4} {5}
                uniaxialMaterial ElasticMultiLinear {6} -strain {7} {8} 0. {9} {10} -stress {11} {12} 0. {13} {14} 
                uniaxialMaterial Parallel {15} {2} {6}
                uniaxialMaterial Series {16} {0} {15}\n""".format(
                100000 + i,
                " ".join([str(p) for p in list(params[4:])]),
                200000 + i,
                params[1],
                params[0] * (1-params[3]),
                params[3],
                300000 + i,
                -params[2] / 2 - 1,
                -params[2] / 2,
                params[2] / 2,
                params[2] / 2 + 1,
                params[3] * params[0] * -params[2] / 2 - (1-params[3]) * params[0],
                params[3] * params[0] * -params[2] / 2,
                params[3] * params[0] * params[2] / 2,
                params[3] * params[0] * params[2] / 2 + (1-params[3]) * params[0],
                400000 + i,
                i
            )
    elif template == "gap_saws":
        dofs = 2
        for i, params in enumerate(paramlist):
            matstring += """uniaxialMaterial SAWS {0} {1}
                uniaxialMaterial ElasticMultiLinear {2} -strain {3} {4} 0. {5} {6} -stress {7} {8} 0. {9} {10}
                """.format(
                i,
                " ".join([str(p) for p in list(params[2:])]),
                100 + i,
                -params[0] - 1,
                -params[0],
                params[0],
                params[0] + 1,
                1.0 / params[1] * params[5] * -params[0] - params[1] * params[5],
                1.0 / params[1] * params[5] * -params[0],
                1.0 / params[1] * params[5] * params[0],
                1.0 / params[1] * params[5] * params[0] + params[1] * params[5],
            )
    elif template == "hook_saws":
        # params: stiffness, activation, slot, sliding_stiffness, SAWS...
        dofs = 2
        for i, params in enumerate(paramlist):
            k0 = params[0] - params[3] / 2.0
            pyratio = params[3] / 2.0 / (params[0] - params[3] / 2.0)
            matstring += """uniaxialMaterial SAWS {0} {1}
            uniaxialMaterial Steel01 {2} {3} {4} {5}
            uniaxialMaterial ElasticMultiLinear {6} -strain {7} {8} 0. {9} {10} -stress {11} {12} 0. {13} {14}
            uniaxialMaterial Parallel {15} {2} {6}
            """.format(
                i,
                " ".join([str(p) for p in list(params[4:])]),
                100000 + i,
                params[1],
                k0,
                pyratio,
                200000 + i,
                -params[2] / 2.0 - 1,
                -params[2] / 2.0,
                params[2] / 2.0,
                params[2] / 2 + 1,
                params[3] / 2.0 * -params[2] / 2.0 - (k0 - params[3] / 2.0),
                params[3] / 2.0 * -params[2] / 2.0,
                params[3] / 2.0 * params[2] / 2.0,
                params[3] / 2.0 * params[2] / 2.0 + (k0 - params[3] / 2.0),
                100 + i,
            )
    elif template == "hook":
        for i, params in enumerate(paramlist):
            k0 = params[0] - params[3] / 2.0
            pyratio = params[3] / 2.0 / (params[0] - params[3] / 2.0)
            matstring += """uniaxialMaterial Steel01 {0} {1} {2} {3}
            uniaxialMaterial ElasticMultiLinear {4} -strain {5} {6} 0. {7} {8} -stress {9} {10} 0. {11} {12}
            uniaxialMaterial Parallel {13} {0} {4}\n""".format(
                100000 + i,
                params[1],
                k0,
                pyratio,
                200000 + i,
                -params[2] / 2.0 - 1,
                -params[2] / 2.0,
                params[2] / 2.0,
                params[2] / 2.0 + 1,
                params[3] / 2.0 * -params[2] / 2.0 - (k0 - params[3] / 2.0),
                params[3] / 2.0 * -params[2] / 2.0,
                params[3] / 2.0 * params[2] / 2.0,
                params[3] / 2.0 * params[2] / 2.0 + (k0 - params[3] / 2.0),
                i
            )
    else:
        for i, params in enumerate(paramlist):
            matstring += "uniaxialMaterial {0} {1} {2}\n".format(
                template,
                i,
                " ".join([str(p) for p in list(params)])
            )
    dispstring = " ".join([str(d) for d in list(disps)])
    script = f"set disps [list {dispstring}]\n set matcount {len(paramlist)}\n set dofs {dofs}\n"
    script += """
    wipe
    wipeAnalysis
    model BasicBuilder -ndm 1 -ndf 1
    node 0 0.
    fix 0 1
    node 1 1.
    """
    script += matstring
    script += """
    for {set i 0} {$i < $matcount} {incr i} {
        if {$dofs == 1} {
            element twoNodeLink $i 0 1 -mat $i -dir 1 -orient 1.0 0.0 0.0 0.0 1.0 0.0
        } else {
            for {set m 1} {$m < $dofs} {incr m} {
                node [expr 100*$m + $i] [expr 1.0 / $dofs]
                element twoNodeLink [expr 100*$m + $i] [expr 100*($m-1) + $i] [expr 100*$m + $i] -mat [expr 100*$m + $i] -dir 1 -orient 1.0 0.0 0.0 0.0 1.0 0.0
            }
            element twoNodeLink $i [expr 100*($dofs-1) + $i] 1 -mat $i -dir 1 -orient 1.0 0.0 0.0 0.0 1.0 0.0
        }
    }
    """
    script += """
    pattern Plain 1 Linear {
        load 1 1.
    }
    constraints Transformation
    numberer Plain
    system BandGeneral
    if {$dofs == 1} {
        test EnergyIncr 1.0e-6 200
    } else {
        test EnergyIncr 1.0e-2 200
    }
    algorithm KrylovNewton
    analysis Static
    set fs [list ]
    set ok 0
    for {set i 0} {$i < [llength $disps]} {incr i} {
        if {$i == 0} {
            integrator DisplacementControl 1 1 [lindex $disps 0]
        } else {
            integrator DisplacementControl 1 1 [expr [lindex $disps $i] - [lindex $disps $i-1]]
        }
        if {$ok == 0} {
            set ok [analyze 1]
        }
        if {$ok == 0} {
            set step [list ]
            for {set j 0} {$j < $matcount} {incr j} {
                lappend step [lindex [eleResponse $j material 1 stress] 0]
            }
        } else {
            set step [list ]
            for {set j 0} {$j < $matcount} {incr j} {
                lappend step 0.0
            }
        }
        lappend fs $step
    }
    puts "FORCERESULT"
    if {[llength $fs] > 0} {
        puts [join $fs " "]
    }
    """
    return script


def run_model(template, paramlist, displist, timeout=10, log=False):
    """Run OpenSees model, and get a ndarray of force results.
    Returns:
        list of ndarray(NDisp, NMat)
    """
    displist = displist if isinstance(displist[0], list) else [displist]
    print(f"Running {template} with {len(paramlist)} mats and {len(displist)} protocols.")
    forcelist = []
    for disps in displist:
        script = generate_script(template, paramlist, disps)
        forces = np.zeros((len(disps), len(paramlist)))
        try:
            res = subprocess.run("OpenSees", universal_newlines=True, 
                                 input=script, stderr=subprocess.PIPE,
                                 timeout=timeout)
            if log:
                print(script)
                print(res.stderr)
            forcestring = res.stderr.split("FORCERESULT\n")[-1]
            if forcestring:
                forces_raw = np.fromstring(forcestring, sep=' ')
                if len(forces_raw) > 1:
                    forces = forces_raw.reshape((len(disps), len(paramlist)))
        except subprocess.TimeoutExpired:
            print("Subprocess time out. Params:")
            if len(paramlist) <= 20:
                for p in paramlist:
                    print(p)
        finally:
            forcelist.append(forces)
    return forcelist


def err_shape(tds, tfs, mfs):
    """Calculate shape error."""
    exts = find_extremes(tds)
    dfs = np.array(tfs).flatten() - np.array(mfs).flatten()
    weights = np.ones(len(tds)) / len(tds)
    for i in range(1, len(exts)):
        if exts[i] - exts[i-1] > 1:
            weights[exts[i-1]+1:exts[i]] /= (exts[i] - exts[i-1] - 1)
    return np.sum(weights * dfs ** 2) 


def cumenergy(ds, fs):
    delta = np.diff(ds)
    fs = np.array(fs)
    average = (fs[:-1] + fs[1:]) / 2
    energy = delta * average
    cumsum = np.cumsum(energy)
    distance = np.cumsum(abs(delta))
    cumsum = np.hstack([0, cumsum])
    distance = np.hstack([0, distance])
    return distance, cumsum


def err_cumenergy(tds, tfs, mfs):
    """Calculate cumulative error for given target disps, target forces, and model forces"""
    delta = np.diff(tds)
    tfs = np.array(tfs)
    mfs = np.array(mfs)
    average_t = (tfs[:-1] + tfs[1:]) / 2
    average_m = (mfs[:-1] + mfs[1:]) / 2
    energy_t = delta * average_t
    energy_m = delta * average_m
    cumsum_t = np.cumsum(energy_t)
    cumsum_m = np.cumsum(energy_m)
    err = abs(cumsum_t - cumsum_m) / abs(cumsum_t)
    distance = np.cumsum(abs(delta))
    err = np.hstack([0, err])
    distance = np.hstack([0, distance])
    return distance, err


def err_cumforce(tds, tfs, mfs):
    delta = np.diff(tds)
    tfs = np.array(tfs)
    mfs = np.array(mfs)
    diff = abs(tfs - mfs)
    distance = np.hstack([0, np.cumsum(abs(delta))])
    return distance, np.cumsum(diff) / np.cumsum(abs(tfs))


class CalibrationProblem(ea.Problem):
    def __init__(self, template, tds, tfs, lower, upper, precisions, groupsize=200):
        """Args: `tds`: list of target disps. `tfs`: list of target forces"""
        name = "Base"
        maxormins = [1]
        Dim = len(precisions) # number of params
        varTypes = np.array([1] * Dim) # integer variables
        self.precisions = precisions
        self.lower = lower
        self.upper = upper
        lb = [0] * Dim
        ub = [np.floor((upper[i]-lower[i]) / precisions[i]) for i in range(Dim)]
        lbin = [1] * Dim
        ubin = [1] * Dim
        self.upperbound = ub
        self.tds = tds
        self.tfs = tfs
        self.tds_all = np.hstack(tds).reshape((-1, 1))  # stacked column vector
        self.tfs_all = np.hstack(tfs).reshape((-1, 1))  # stacked column vector
        weightslist = []
        for n in range(len(tds)):
            weights = np.ones(len(self.tfs[n])) / len(self.tfs[n])
            exts = find_extremes(self.tds[n])
            for i in range(1, len(exts)):
                if exts[i] - exts[i-1] - 1 > 0:
                    weights[exts[i-1]+1:exts[i]] /= (exts[i] - exts[i-1] - 1)
            weightslist.append(weights)
        self.weights = weightslist
        self.weights_all = np.hstack(weightslist).reshape((-1, 1))
        self.groupsize = groupsize
        self.template = template
        ea.Problem.__init__(self, name, 1, maxormins, Dim, varTypes, lb, ub, lbin, ubin)

    def aimFunc(self, pop):
        vs = pop.Phen
        forces = np.zeros((len(self.tfs_all), vs.shape[0]))
        ps = vs * self.precisions + self.lower
        allparams = [list(ps[i, :]) for i in range(vs.shape[0])]
        for i in range(len(allparams) // self.groupsize):
            start = i * self.groupsize
            end = (i+1) * self.groupsize
            paramlist = allparams[start:end]
            if not paramlist:
                continue
            forcelist = run_model(self.template, paramlist, self.tds, timeout=30, log=False)
            forces[:, start:end] = np.vstack(forcelist)
        err = 0
        i = 0
        squares = self.weights_all * (forces - self.tfs_all) ** 2
        err = np.sum(squares, axis=0)
        pop.ObjV = err.reshape((-1, 1))

    def calReferObjV(self):
        return np.array([[0.]])


def runGA(name, template, tds, tfs, precisions, lb, ub, groupsize=200, NIND=1000, GEN=10):
    problem = CalibrationProblem(template, tds, tfs, lb, ub, precisions, groupsize)
    problem.name = name
    Encoding = 'RI'
    NIND = NIND
    Field = ea.crtfld(Encoding, problem.varTypes, problem.ranges, problem.borders)
    population = ea.Population(Encoding, Field, NIND)
    myAlgorithm = ea.soea_SEGA_templet(problem, population)
    myAlgorithm.MAXGEN = GEN
    myAlgorithm.drawing = 2
    [population, obj_trace, var_trace] = myAlgorithm.run()
    #output
    best_gen = np.argmin(problem.maxormins * obj_trace[:, 1])
    best_ObjV = obj_trace[best_gen, 1]
    print('最优的目标函数值为：%s'%(best_ObjV))
    print('最优的控制变量值为：')
    for i in range(var_trace.shape[1]):
        print(i, var_trace[best_gen, i], "max", problem.upperbound[i])
        if (not 0.05 < var_trace[best_gen, i] / problem.upperbound[i] < 0.95):
            print("WARNING")
    print('有效进化代数：%s'%(obj_trace.shape[0]))
    print('最优的一代是第 %s 代'%(best_gen + 1))
    print('评价次数：%s'%(myAlgorithm.evalsNum))
    print('时间已过 %s 秒'%(myAlgorithm.passTime))
    # save
    with open(f'optimal/{name}', 'w') as f:
        var = var_trace[best_gen, :]
        opt = var * np.array(precisions) + np.array(problem.lower)
        print("Optimal params:", opt)
        f.write("\n".join([str(o) for o in opt]))
        f.write("\n")
        f.write(str(best_ObjV))